Example 2: Understanding the inputs and the output trace
Contents
Example 2: Understanding the inputs and the output trace¶
In this notebook, we will discuss more in detail the inputs and the output trace of the VBMC algorithm, including a discussion on how to set parameter bounds and starting point. We also describe the items in the output trace, and display the evolution of the variational posterior in each iteration.
This notebook is part 1 of a series of notebooks in which we present various example usages for VBMC with the pyvbmc package.
import numpy as np
import scipy.stats as scs
from pyvbmc.vbmc import VBMC
1. Model definition: log likelihood, log prior, and log joint¶
We use the same toy target function as in Example 1, a broad Rosenbrock’s banana function in \(D = 2\).
D = 2 # still in 2-D
# log likelihood (Rosenbrock)
def llfun(x):
x = np.atleast_2d(x)
return -np.sum(
(x[0, :-1] ** 2.0 - x[0, 1:]) ** 2.0 + (x[0, :-1] - 1) ** 2.0 / 100
)
To make the problem more interesting, we will assume parameters to be constrained to be positive, that is \(x_d > 0\) for \(1 \le d \le D\).
Since parameters are strictly positive, we impose an independent exponential prior on each variable (again, you could choose here whatever you want). We then define the log joint (our target) as log likelihood plus log prior.
prior_tau = 3 * np.ones((1, D)) # Length scale of the exponential prior
lpriorfun = lambda x: np.sum(scs.expon.logpdf(x, scale=prior_tau))
target = lambda x: llfun(x) + lpriorfun(x) # Log joint
2. Parameter setup: bounds and starting point¶
Choosing parameter bounds¶
Bound settings require some discussion:
the specified (hard) bounds are not included in the domain, so in this case, since we want to be the parameters to be positive, we can set the lower bounds
LB = 0, knowing that the parameter will always be greater than zero.
LB = np.zeros((1, D)) # Lower bounds
Currently, VBMC does not support half-bounds, so we need to specify a finite upper bound (cannot be
inf). We pick something very large according to our prior. However, do not go crazy here by picking something impossibly large, otherwise VBMC will likely fail.
UB = 10 * prior_tau # Upper bounds
Plausible bounds need to be meaningfully different from the hard bounds, so here we cannot set the plausible lower bounds
PLB = 0. We follow the same strategy as the first example, by picking the top ~68% credible interval of the prior.
PLB = scs.expon.ppf(0.159, scale=prior_tau)
PUB = scs.expon.ppf(0.841, scale=prior_tau)
Choosing the starting point¶
Good practice is to initialize VBMC in a region of high posterior density.
With no additional information, uninformed approaches to choose the starting point \(\mathbf{x}_0\) include picking the mode of the prior, or, if you need multiple points, a random draw from the prior or a random point uniformly chosen inside the plausible box. None of these approaches is particularly recommended, though, especially when facing more difficult problems.
It is generally better to adopt an informed starting point \(\mathbf{x}_0\). A good choice is typically a point in the approximate vicinity of the mode of the posterior (no need to start exactly at the mode). Clearly, it is hard to know in advance the location of the mode, but a common approach is to run a few optimization steps on the target (starting from one of the uninformed options above).
For this example, we cheat and start in the proximity of the mode, exploiting our knowledge of the toy likelihood, for which we know the location of the maximum.
x0 = np.ones((1, D)) # Optimum of the Rosenbrock function
3. Initialize and run inference¶
Finally, we initialize and run VBMC.
Output trace¶
In each iteration, the VBMC output trace prints the following information:
Iteration: iteration number (starting from 0, which follows the initial design).f-count: total number of thetargetfunction evaluations so far.Mean[ELBO]: Estimated mean of theelbo.Std[ELBO]: Standard deviation of the estimatedelbo.sKL-iter[q]: Gaussianized symmetrized Kullback-Leibler divergence between the variational posterior in the previous and current iteration. In brief, a measure of how much the variational posterior \(q\) changes between iterations.K[q]: Number of mixture components of the variational posterior.Convergence: A number that summarizes the stability of the current solution, by weighing various factors such as the change in theelboand in the variational posterior across iterations, and the uncertainty in theelbo. Values (much) less than 1 correspond to a stable solution and thus suggest convergence of the algorithm.Action: Major actions taken by the algorithm in the current iteration. Common actions/events are:start warm-upandend warm-up: start and end of the initial warm-up stage;trim data: some points are removed to increase stability;stable: solution stable for multiple iteration;finalize: refine the variational approximation with additional mixture components.
Iteration plotting¶
In this example, we also switch on the plotting of the current variational posterior at each iteration with the user option plot. Regarding the iteration plots:
Each iteration plot shows the multidimensional variational posterior in the current iteration via one- and two-dimensional marginal plots for each variable and pairs of variables, as you would get from
vp.plot()(see also Example 1).In each plot, blue circles represent points in the training set, and orange dots are points added to the training set in this iteration.
The red crosses are the centers of the variational mixture components.
Since VBMC with plotting on can occupy a lot of space on the notebook, remember that in Jupyter notebook you can toggle the scrolling on and off via the “Cell > Current Outputs > Toggle Scrolling” and “Cell > All Output > Toggle Scrolling” menus.
# initialize VBMC
user_options = {"plot": True}
vbmc = VBMC(target, x0, LB, UB, PLB, PUB, user_options)
# run VBMC
vp, elbo, elbo_sd, _, _ = vbmc.optimize()
Beginning variational optimization assuming EXACT observations of the log-joint.
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
0 10 -3.20 1.46 191163.68 2 inf start warm-up
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
1 15 -2.93 0.15 0.17 2 inf
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
2 20 -2.61 0.17 0.09 2 3.74
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
3 25 -2.53 0.03 0.06 2 1.8
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
4 30 -2.09 1.99 0.53 2 20.6 end warm-up
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
5 35 -2.40 0.05 0.14 2 4.38
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
6 40 -2.38 0.05 0.04 2 1.18
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
7 45 -2.19 0.04 0.07 3 2.51
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
8 50 -2.17 0.01 0.02 4 0.473
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
9 55 -2.14 0.01 0.02 7 0.57
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
10 60 -2.09 0.01 0.01 10 0.377
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
11 65 -2.05 0.01 0.01 13 0.328
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
12 70 -2.01 0.03 0.01 14 0.375
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
13 75 -2.02 0.00 0.00 14 0.183
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
14 80 -1.99 0.00 0.00 15 0.116
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
15 85 -2.01 0.00 0.00 16 0.0625
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
16 90 -1.99 0.00 0.00 15 0.0762
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
17 95 -1.98 0.00 0.00 14 0.046 stable
Iteration f-count Mean[ELBO] Std[ELBO] sKL-iter[q] K[q] Convergence Action
inf 95 -1.93 0.01 0.00 50 0.046 finalize
Inference terminated: variational solution stable for options.tolstablecountfcn evaluations.
Estimated ELBO: -1.926 +/-0.006.
lml_true = -1.836 # ground truth, which we know for this toy scenario
print("The true log model evidence is:", lml_true)
print("The obtained ELBO is:", format(elbo, ".3f"))
print("The obtained ELBO_SD is:", format(elbo_sd, ".3f"))
The true log model evidence is: -1.836
The obtained ELBO is: -1.926
The obtained ELBO_SD is: 0.006
4. Examine and visualize results¶
We will examine now the obtained variational posterior with an interactive plot.
Note: The interactive plot will not be visible in a notebook preview on GitHub. See below for a static image.
# Generate samples from the variational posterior
n_samples = int(3e5)
Xs, _ = vp.sample(n_samples)
# We compute the pdf of the approximate posterior on a 2-D grid
plot_lb = np.zeros(2)
plot_ub = np.quantile(Xs, 0.999, axis=0)
x1 = np.linspace(plot_lb[0], plot_ub[0], 400)
x2 = np.linspace(plot_lb[1], plot_ub[1], 400)
xa, xb = np.meshgrid(x1, x2) # Build the grid
xx = np.vstack(
(xa.ravel(), xb.ravel())
).T # Convert grids to a vertical array of 2-D points
yy = vp.pdf(xx) # Compute PDF values on specified points
# Plot approximate posterior pdf (this interactive plot does not work in higher D)
import plotly.graph_objects as go
fig = go.Figure(
data=[
go.Surface(z=yy.reshape(x1.size, x2.size), x=x1, y=x2, showscale=False)
]
)
fig.update_layout(
autosize=False,
width=800,
height=500,
scene=dict(
xaxis_title="x_0",
yaxis_title="x_1",
zaxis_title="Approximate posterior pdf",
),
)
# Compute and plot approximate posterior mean
post_mean = np.mean(Xs, axis=0)
fig.add_trace(
go.Scatter3d(
x=[post_mean[0]],
y=[post_mean[1]],
z=vp.pdf(post_mean),
name="Approximate posterior mean",
marker_symbol="diamond",
)
)
# Find and plot approximate posterior mode
# post_mode = vp.mode()
# fig.add_trace(go.Scatter3d(x=[post_mode[0]], y=[post_mode[1]], z=vp.pdf(post_mode), name="Approximate posterior mode", marker_symbol="x"))
# Display posterior also as a non-interactive static image
from IPython.display import Image
Image(fig.to_image(format="png"))
5. Conclusions¶
In this notebook, we have discussed more in detail the input parameter settings (bounds and starting point) and the output trace of VBMC.
An important step still not considered in this example is a thorough validation of the results and diagnostics, which will be discussed in a later notebook.
Example 2: full code¶
The following cell includes in a single place all the code used in Example 2, without the extra fluff.
import numpy as np
import scipy.stats as scs
from pyvbmc.vbmc import VBMC
#####################################################################
# 1. Model definition: log likelihood, log prior, and log joint
D = 2 # still in 2-D
# log likelihood (Rosenbrock)
def llfun(x):
x = np.atleast_2d(x)
return -np.sum(
(x[0, :-1] ** 2.0 - x[0, 1:]) ** 2.0 + (x[0, :-1] - 1) ** 2.0 / 100
)
prior_tau = 3 * np.ones((1, D)) # Length scale of the exponential prior
lpriorfun = lambda x: np.sum(scs.expon.logpdf(x, scale=prior_tau))
target = lambda x: llfun(x) + lpriorfun(x) # Log joint
#####################################################################
# 2. Parameter setup: bounds and starting point
LB = np.zeros((1, D)) # Lower bounds
UB = 10 * prior_tau # Upper bounds
PLB = scs.expon.ppf(0.159, scale=prior_tau)
PUB = scs.expon.ppf(0.841, scale=prior_tau)
x0 = np.ones((1, D)) # Optimum of the Rosenbrock function
#####################################################################
# 3. Initialize and run inference
# initialize VBMC
user_options = {"plot": True}
vbmc = VBMC(target, x0, LB, UB, PLB, PUB, user_options)
# run VBMC
vp, elbo, elbo_sd, _, _ = vbmc.optimize()
lml_true = -1.836 # ground truth, which we know for this toy scenario
print("The true log model evidence is:", lml_true)
print("The obtained ELBO is:", format(elbo, ".3f"))
print("The obtained ELBO_SD is:", format(elbo_sd, ".3f"))
#####################################################################
# 4. Analyze and visualize results
# Generate samples from the variational posterior
n_samples = int(3e5)
Xs, _ = vp.sample(n_samples)
# We compute the pdf of the approximate posterior on a 2-D grid
plot_lb = np.zeros(2)
plot_ub = np.quantile(Xs, 0.999, axis=0)
x1 = np.linspace(plot_lb[0], plot_ub[0], 400)
x2 = np.linspace(plot_lb[1], plot_ub[1], 400)
xa, xb = np.meshgrid(x1, x2) # Build the grid
xx = np.vstack(
(xa.ravel(), xb.ravel())
).T # Convert grids to a vertical array of 2-D points
yy = vp.pdf(xx) # Compute PDF values on specified points
# Plot approximate posterior pdf (this interactive plot does not work in higher D)
import plotly.graph_objects as go
fig = go.Figure(
data=[
go.Surface(z=yy.reshape(x1.size, x2.size), x=x1, y=x2, showscale=False)
]
)
fig.update_layout(
autosize=False,
width=800,
height=500,
scene=dict(
xaxis_title="x_0",
yaxis_title="x_1",
zaxis_title="Approximate posterior pdf",
),
)
# Compute and plot approximate posterior mean
post_mean = np.mean(Xs, axis=0)
fig.add_trace(
go.Scatter3d(
x=[post_mean[0]],
y=[post_mean[1]],
z=vp.pdf(post_mean),
name="Approximate posterior mean",
marker_symbol="diamond",
)
)
# Find and plot approximate posterior mode
# post_mode = vp.mode()
# fig.add_trace(go.Scatter3d(x=[post_mode[0]], y=[post_mode[1]], z=vp.pdf(post_mode), name="Approximate posterior mode", marker_symbol="x"))
# Display posterior also as a non-interactive static image
from IPython.display import Image
Image(fig.to_image(format="png"))
Acknowledgments¶
Work on the pyvbmc package was funded by the Finnish Center for Artificial Intelligence FCAI.